// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <memory>
#include <string>
#include <vector>

#include <ktest/ktest.h>

#include <nebula/compute/function.h>
#include <nebula/compute/function_options.h>
#include <nebula/compute/registry.h>

#include <turbo/utility/status.h>
#include <nebula/testing/ktest_util.h>
#include <turbo/base/macros.h>

namespace nebula::compute {

    using MakeFunctionRegistry = std::function<std::unique_ptr<FunctionRegistry>()>;
    using GetNumFunctions = std::function<int()>;
    using get_function_names = std::function<std::vector<std::string>()>;
    using TestRegistryParams =
            std::tuple<MakeFunctionRegistry, GetNumFunctions, get_function_names, std::string>;

    struct TestRegistry : public ::testing::TestWithParam<TestRegistryParams> {
    };

    TEST(TestRegistry, create_builtin_registry) {
        // This does KCHECK_OK internally for now so this will fail in debug builds
        // if there is a problem initializing the global function registry
        FunctionRegistry *registry = get_function_registry();
        TURBO_UNUSED(registry);
    }

    TEST_P(TestRegistry, Basics) {
        auto registry_factory = std::get<0>(GetParam());
        auto registry_ = registry_factory();
        auto get_num_funcs = std::get<1>(GetParam());
        int n_funcs = get_num_funcs();
        auto get_func_names = std::get<2>(GetParam());
        std::vector<std::string> func_names = get_func_names();
        ASSERT_EQ(n_funcs, registry_->num_functions());

        std::shared_ptr<Function> func = std::make_shared<ScalarFunction>(
                "f1", Arity::Unary(), /*doc=*/FunctionDoc::Empty());
        ASSERT_OK(registry_->add_function(func));
        ASSERT_EQ(n_funcs + 1, registry_->num_functions());

        func = std::make_shared<VectorFunction>("f0", Arity::Binary(),
                /*doc=*/FunctionDoc::Empty());
        ASSERT_OK(registry_->add_function(func));
        ASSERT_EQ(n_funcs + 2, registry_->num_functions());

        ASSERT_OK_AND_ASSIGN(std::shared_ptr<const Function> f1, registry_->get_function("f1"));
        ASSERT_EQ("f1", f1->name());

        // Nonexistent function
        ASSERT_RAISES(turbo::StatusCode::kNotFound, registry_->get_function("f2"));

        // Try adding a function with name collision
        func = std::make_shared<ScalarAggregateFunction>("f1", Arity::Unary(),
                /*doc=*/FunctionDoc::Empty());
        ASSERT_RAISES(turbo::StatusCode::kAlreadyExists, registry_->add_function(func));

        // Allow overwriting by flag
        ASSERT_OK(registry_->add_function(func, /*allow_overwrite=*/true));
        ASSERT_OK_AND_ASSIGN(f1, registry_->get_function("f1"));
        ASSERT_EQ(Function::SCALAR_AGGREGATE, f1->kind());

        std::vector<std::string> expected_names(func_names);
        for (auto name: {"f0", "f1"}) {
            expected_names.push_back(name);
        }
        std::sort(expected_names.begin(), expected_names.end());
        ASSERT_EQ(expected_names, registry_->get_function_names());

        // Aliases
        ASSERT_RAISES(turbo::StatusCode::kNotFound, registry_->add_alias("f33", "f3"));
        ASSERT_OK(registry_->add_alias("f11", "f1"));
        ASSERT_OK_AND_ASSIGN(std::shared_ptr<const Function> f2, registry_->get_function("f11"));
        ASSERT_EQ(func, f2);
    }

// Define a custom print since the default Googletest print trips Valgrind
    void print_to(const TestRegistryParams &param, std::ostream *os) {
        (*os) << "TestRegistryParams{"
              << "get_num_funcs()=" << std::get<1>(param)() << ", get_func_names()=";
        for (std::string func_name: std::get<2>(param)()) {
            (*os) << func_name;
        }
        (*os) << ", name=" << std::get<3>(param) << "}";
    }

    INSTANTIATE_TEST_SUITE_P(
            TestRegistry, TestRegistry,
            testing::Values(
                    TestRegistryParams{std::make_tuple(
                            static_cast<MakeFunctionRegistry>([]() { return FunctionRegistry::create(); }),
                            []() { return 0; }, []() { return std::vector<std::string>{}; }, "default")},
                    TestRegistryParams{std::make_tuple(
                            static_cast<MakeFunctionRegistry>([]() {
                                return FunctionRegistry::create(get_function_registry());
                            }),
                            []() { return get_function_registry()->num_functions(); },
                            []() { return get_function_registry()->get_function_names(); }, "nested")}));

    TEST(TestRegistry, RegisterTempFunctions) {
        auto default_registry = get_function_registry();
        constexpr int rounds = 3;
        for (int i = 0; i < rounds; i++) {
            auto registry = FunctionRegistry::create(default_registry);
            for (std::string func_name: {"f1", "f2"}) {
                std::shared_ptr<Function> func = std::make_shared<ScalarFunction>(
                        func_name, Arity::Unary(), /*doc=*/FunctionDoc::Empty());
                ASSERT_OK(registry->can_add_function(func));
                ASSERT_OK(registry->add_function(func));
                ASSERT_RAISES(turbo::StatusCode::kAlreadyExists, registry->can_add_function(func));
                ASSERT_RAISES(turbo::StatusCode::kAlreadyExists, registry->add_function(func));
                ASSERT_OK(default_registry->can_add_function(func));
            }
        }
    }

    TEST(TestRegistry, RegisterTempAliases) {
        auto default_registry = get_function_registry();
        std::vector<std::string> func_names = default_registry->get_function_names();
        constexpr int rounds = 3;
        for (int i = 0; i < rounds; i++) {
            auto registry = FunctionRegistry::create(default_registry);
            for (std::string func_name: func_names) {
                std::string alias_name = "alias_of_" + func_name;
                std::shared_ptr<Function> func = std::make_shared<ScalarFunction>(
                        func_name, Arity::Unary(), /*doc=*/FunctionDoc::Empty());
                ASSERT_RAISES(turbo::StatusCode::kNotFound, registry->get_function(alias_name));
                ASSERT_OK(registry->can_add_alias(alias_name, func_name));
                ASSERT_OK(registry->add_alias(alias_name, func_name));
                ASSERT_OK(registry->get_function(alias_name));
                ASSERT_OK(default_registry->get_function(func_name));
                ASSERT_RAISES(turbo::StatusCode::kNotFound, default_registry->get_function(alias_name));
            }
        }
    }

    template<int kExampleSeqNum>
    class ExampleOptions : public FunctionOptions {
    public:
        explicit ExampleOptions(std::shared_ptr<Scalar> value);

        std::shared_ptr<Scalar> value;
    };

    template<int kExampleSeqNum>
    class ExampleOptionsType : public FunctionOptionsType {
    public:
        static const FunctionOptionsType *GetInstance() {
            static std::unique_ptr<FunctionOptionsType> instance(
                    new ExampleOptionsType<kExampleSeqNum>());
            return instance.get();
        }

        const char *type_name() const override {
            static std::string name = std::string("example") + std::to_string(kExampleSeqNum);
            return name.c_str();
        }

        std::string stringify(const FunctionOptions &options) const override {
            return type_name();
        }

        bool compare(const FunctionOptions &options,
                     const FunctionOptions &other) const override {
            return true;
        }

        std::unique_ptr<FunctionOptions> copy(const FunctionOptions &options) const override {
            const auto &opts = static_cast<const ExampleOptions<kExampleSeqNum> &>(options);
            return std::make_unique<ExampleOptions<kExampleSeqNum>>(opts.value);
        }
    };

    template<int kExampleSeqNum>
    ExampleOptions<kExampleSeqNum>::ExampleOptions(std::shared_ptr<Scalar> value)
            : FunctionOptions(ExampleOptionsType<kExampleSeqNum>::GetInstance()),
              value(std::move(value)) {}

    TEST(TestRegistry, RegisterTempFunctionOptionsType) {
        auto default_registry = get_function_registry();
        std::vector<const FunctionOptionsType *> options_types = {
                ExampleOptionsType<1>::GetInstance(),
                ExampleOptionsType<2>::GetInstance(),
        };
        constexpr int rounds = 3;
        for (int i = 0; i < rounds; i++) {
            auto registry = FunctionRegistry::create(default_registry);
            for (auto options_type: options_types) {
                ASSERT_OK(registry->can_add_function_options_type(options_type));
                ASSERT_OK(registry->add_function_options_type(options_type));
                ASSERT_RAISES(turbo::StatusCode::kAlreadyExists, registry->can_add_function_options_type(options_type));
                ASSERT_RAISES(turbo::StatusCode::kAlreadyExists, registry->add_function_options_type(options_type));
                ASSERT_OK(default_registry->can_add_function_options_type(options_type));
            }
        }
    }

    TEST(TestRegistry, RegisterNestedFunctions) {
        auto default_registry = get_function_registry();
        std::shared_ptr<Function> func1 = std::make_shared<ScalarFunction>(
                "f1", Arity::Unary(), /*doc=*/FunctionDoc::Empty());
        std::shared_ptr<Function> func2 = std::make_shared<ScalarFunction>(
                "f2", Arity::Unary(), /*doc=*/FunctionDoc::Empty());
        constexpr int rounds = 3;
        for (int i = 0; i < rounds; i++) {
            auto registry1 = FunctionRegistry::create(default_registry);

            ASSERT_OK(registry1->can_add_function(func1));
            ASSERT_OK(registry1->add_function(func1));

            for (int j = 0; j < rounds; j++) {
                auto registry2 = FunctionRegistry::create(registry1.get());

                ASSERT_OK(registry2->can_add_function(func2));
                ASSERT_OK(registry2->add_function(func2));
                ASSERT_RAISES(turbo::StatusCode::kAlreadyExists, registry2->can_add_function(func2));
                ASSERT_RAISES(turbo::StatusCode::kAlreadyExists, registry2->add_function(func2));
                ASSERT_OK(default_registry->can_add_function(func2));
            }

            ASSERT_RAISES(turbo::StatusCode::kAlreadyExists, registry1->can_add_function(func1));
            ASSERT_RAISES(turbo::StatusCode::kAlreadyExists, registry1->add_function(func1));
            ASSERT_OK(default_registry->can_add_function(func1));
        }
    }

}  // namespace nebula::compute
